import re
import json
from typing import List, Dict, Any, Optional
from datetime import datetime

from biomni.llm import get_llm
from logger_config import get_logger
from config import Config

logger = get_logger(__name__)

class SummaryService:
    """智能总结服务，使用LLM生成执行过程的智能总结"""
    
    def __init__(self):
        self.config = Config()
        self.llm = None
        self._initialize_llm()
    
    def _initialize_llm(self):
        """初始化LLM实例"""
        try:
            self.llm = get_llm(
                model=self.config.DEFAULT_MODEL,
                temperature=0.3  # 使用较低的温度以获得更一致的总结
            )
            logger.info(f"智能总结服务LLM初始化成功，模型: {self.config.DEFAULT_MODEL}")
        except Exception as e:
            logger.error(f"智能总结服务LLM初始化失败: {str(e)}")
            self.llm = None
    
    def generate_intelligent_summary(self, 
                                   execution_log: List[Dict[str, Any]], 
                                   query: str, 
                                   execution_time: float,
                                   final_result: Any = None) -> Dict[str, Any]:
        """生成智能执行总结
        
        Args:
            execution_log: 执行日志列表
            query: 用户查询
            execution_time: 执行时间
            final_result: 最终执行结果
            
        Returns:
            智能总结字典
        """
        try:
            if not self.llm:
                logger.warning("LLM未初始化，使用基础总结")
                return self._generate_basic_summary(execution_log, query, execution_time)
            
            # 检查执行日志长度，决定是否需要分段总结
            if self._should_use_segmented_summary(execution_log):
                return self._generate_segmented_summary(execution_log, query, execution_time, final_result)
            else:
                return self._generate_single_summary(execution_log, query, execution_time, final_result)
                
        except Exception as e:
            logger.error(f"生成智能总结失败: {str(e)}")
            return self._generate_basic_summary(execution_log, query, execution_time)
    
    def _should_use_segmented_summary(self, execution_log: List[Dict[str, Any]]) -> bool:
        """判断是否需要使用分段总结"""
        # 计算总的文本长度
        total_length = sum(len(str(log_item)) for log_item in execution_log)
        # 如果总长度超过8000字符或日志条目超过50条，使用分段总结
        return total_length > 8000 or len(execution_log) > 50
    
    def _generate_single_summary(self, 
                               execution_log: List[Dict[str, Any]], 
                               query: str, 
                               execution_time: float,
                               final_result: Any = None) -> Dict[str, Any]:
        """生成单次总结"""
        try:
            # 构建提示词
            prompt = self._build_summary_prompt(execution_log, query, execution_time, final_result)
            
            # 调用LLM生成总结
            response = self.llm.invoke(prompt)
            summary_text = response.content if hasattr(response, 'content') else str(response)
            
            # 解析总结结果
            parsed_summary = self._parse_llm_summary(summary_text)
            
            return {
                'type': 'intelligent_summary',
                'timestamp': datetime.now().isoformat(),
                'query': query[:100] + '...' if len(query) > 100 else query,
                'execution_time': f"{execution_time:.2f}s",
                'summary': parsed_summary,
                'method': 'single_pass'
            }
            
        except Exception as e:
            logger.error(f"单次总结生成失败: {str(e)}")
            raise
    
    def _generate_segmented_summary(self, 
                                  execution_log: List[Dict[str, Any]], 
                                  query: str, 
                                  execution_time: float,
                                  final_result: Any = None) -> Dict[str, Any]:
        """生成分段总结"""
        try:
            # 将执行日志分段
            segments = self._segment_execution_log(execution_log)
            segment_summaries = []
            
            # 为每个段生成总结
            for i, segment in enumerate(segments):
                segment_prompt = self._build_segment_summary_prompt(segment, i + 1, len(segments))
                response = self.llm.invoke(segment_prompt)
                segment_summary = response.content if hasattr(response, 'content') else str(response)
                segment_summaries.append(segment_summary)
            
            # 生成最终汇总
            final_prompt = self._build_final_summary_prompt(segment_summaries, query, execution_time, final_result)
            final_response = self.llm.invoke(final_prompt)
            final_summary_text = final_response.content if hasattr(final_response, 'content') else str(final_response)
            
            # 解析最终总结
            parsed_summary = self._parse_llm_summary(final_summary_text)
            
            return {
                'type': 'intelligent_summary',
                'timestamp': datetime.now().isoformat(),
                'query': query[:100] + '...' if len(query) > 100 else query,
                'execution_time': f"{execution_time:.2f}s",
                'summary': parsed_summary,
                'method': 'segmented',
                'segments_count': len(segments)
            }
            
        except Exception as e:
            logger.error(f"分段总结生成失败: {str(e)}")
            raise
    
    def _segment_execution_log(self, execution_log: List[Dict[str, Any]]) -> List[List[Dict[str, Any]]]:
        """将执行日志分段"""
        segments = []
        current_segment = []
        current_length = 0
        max_segment_length = 6000  # 每段最大字符数
        
        for log_item in execution_log:
            item_length = len(str(log_item))
            
            if current_length + item_length > max_segment_length and current_segment:
                segments.append(current_segment)
                current_segment = [log_item]
                current_length = item_length
            else:
                current_segment.append(log_item)
                current_length += item_length
        
        if current_segment:
            segments.append(current_segment)
        
        return segments
    
    def _build_summary_prompt(self, 
                            execution_log: List[Dict[str, Any]], 
                            query: str, 
                            execution_time: float,
                            final_result: Any = None) -> str:
        """构建总结提示词"""
        log_text = self._format_execution_log(execution_log)
        result_text = self._format_final_result(final_result) if final_result else "无最终结果"
        
        prompt = f"""请为以下生物医学AI代理的执行过程生成一个简洁而全面的中文总结。

用户查询：{query}
执行时间：{execution_time:.2f}秒

执行过程：
{log_text}

最终结果：
{result_text}

请生成一个结构化的总结，包含以下内容：
1. 任务概述：简要描述用户的需求和代理要解决的问题
2. 执行步骤：列出主要的执行步骤和关键操作
3. 关键发现：提取重要的数据、结果或洞察
4. 执行状态：说明任务是否成功完成，如有问题请指出
5. 总结建议：基于执行结果给出简要的结论或建议

请用清晰、专业的中文回答，避免过于技术性的细节。总结应该在300-500字之间。"""
        
        return prompt
    
    def _build_segment_summary_prompt(self, segment: List[Dict[str, Any]], segment_num: int, total_segments: int) -> str:
        """构建段落总结提示词"""
        log_text = self._format_execution_log(segment)
        
        prompt = f"""请为以下执行过程片段生成简洁总结（这是第{segment_num}段，共{total_segments}段）：

执行过程片段：
{log_text}

请简要总结这个片段中的主要操作和结果，重点关注：
- 执行了什么操作
- 获得了什么结果
- 是否有错误或问题

总结应该在100-200字之间，用中文回答。"""
        
        return prompt
    
    def _build_final_summary_prompt(self, 
                                  segment_summaries: List[str], 
                                  query: str, 
                                  execution_time: float,
                                  final_result: Any = None) -> str:
        """构建最终汇总提示词"""
        summaries_text = "\n\n".join([f"段落{i+1}总结：{summary}" for i, summary in enumerate(segment_summaries)])
        result_text = self._format_final_result(final_result) if final_result else "无最终结果"
        
        prompt = f"""基于以下各段落的总结，请生成一个完整的执行总结：

用户查询：{query}
执行时间：{execution_time:.2f}秒

各段落总结：
{summaries_text}

最终结果：
{result_text}

请生成一个结构化的完整总结，包含：
1. 任务概述：简要描述用户的需求
2. 执行步骤：整合各段落的主要操作
3. 关键发现：提取重要的结果和洞察
4. 执行状态：整体执行情况
5. 总结建议：基于完整执行过程的结论

请用清晰、专业的中文回答，总结应该在400-600字之间。"""
        
        return prompt
    
    def _format_execution_log(self, execution_log: List[Dict[str, Any]]) -> str:
        """格式化执行日志为文本"""
        formatted_lines = []
        
        for i, log_item in enumerate(execution_log, 1):
            if isinstance(log_item, dict):
                log_type = log_item.get('type', 'unknown')
                content = log_item.get('content', str(log_item))
                
                if log_type == 'reasoning':
                    formatted_lines.append(f"{i}. [思考] {content[:200]}..." if len(content) > 200 else f"{i}. [思考] {content}")
                elif log_type == 'execution':
                    formatted_lines.append(f"{i}. [执行] {content[:200]}..." if len(content) > 200 else f"{i}. [执行] {content}")
                elif log_type == 'observation':
                    formatted_lines.append(f"{i}. [观察] {content[:200]}..." if len(content) > 200 else f"{i}. [观察] {content}")
                elif log_type == 'solution':
                    formatted_lines.append(f"{i}. [解决方案] {content[:200]}..." if len(content) > 200 else f"{i}. [解决方案] {content}")
                else:
                    formatted_lines.append(f"{i}. [日志] {content[:200]}..." if len(content) > 200 else f"{i}. [日志] {content}")
            else:
                content = str(log_item)
                formatted_lines.append(f"{i}. {content[:200]}..." if len(content) > 200 else f"{i}. {content}")
        
        return "\n".join(formatted_lines)
    
    def _format_final_result(self, final_result: Any) -> str:
        """格式化最终结果"""
        if final_result is None:
            return "无最终结果"
        
        if isinstance(final_result, dict):
            # 如果是字典，尝试提取关键信息
            if 'solution' in final_result:
                return f"解决方案：{str(final_result['solution'])[:300]}..."
            elif 'result' in final_result:
                return f"结果：{str(final_result['result'])[:300]}..."
            else:
                return f"结果：{str(final_result)[:300]}..."
        else:
            result_str = str(final_result)
            return result_str[:300] + "..." if len(result_str) > 300 else result_str
    
    def _parse_llm_summary(self, summary_text: str) -> Dict[str, str]:
        """解析LLM生成的总结文本"""
        try:
            # 尝试解析结构化总结
            sections = {
                'overview': '',
                'steps': '',
                'findings': '',
                'status': '',
                'recommendations': ''
            }
            
            # 使用正则表达式提取各个部分
            patterns = {
                'overview': r'(?:任务概述|概述)[：:](.*?)(?=\n\d+\.|\n[^\n]*[：:]|$)',
                'steps': r'(?:执行步骤|步骤)[：:](.*?)(?=\n\d+\.|\n[^\n]*[：:]|$)',
                'findings': r'(?:关键发现|发现)[：:](.*?)(?=\n\d+\.|\n[^\n]*[：:]|$)',
                'status': r'(?:执行状态|状态)[：:](.*?)(?=\n\d+\.|\n[^\n]*[：:]|$)',
                'recommendations': r'(?:总结建议|建议|结论)[：:](.*?)(?=\n\d+\.|\n[^\n]*[：:]|$)'
            }
            
            for key, pattern in patterns.items():
                match = re.search(pattern, summary_text, re.DOTALL | re.IGNORECASE)
                if match:
                    sections[key] = match.group(1).strip()
            
            # 如果没有找到结构化内容，将整个文本作为概述
            if not any(sections.values()):
                sections['overview'] = summary_text.strip()
            
            return sections
            
        except Exception as e:
            logger.error(f"解析LLM总结失败: {str(e)}")
            return {'overview': summary_text.strip()}
    
    def _generate_basic_summary(self, 
                              execution_log: List[Dict[str, Any]], 
                              query: str, 
                              execution_time: float) -> Dict[str, Any]:
        """生成基础总结（当LLM不可用时的后备方案）"""
        try:
            total_steps = len(execution_log)
            step_types = {}
            
            for log_item in execution_log:
                if isinstance(log_item, dict):
                    step_type = log_item.get('type', 'unknown')
                    step_types[step_type] = step_types.get(step_type, 0) + 1
            
            return {
                'type': 'basic_summary',
                'timestamp': datetime.now().isoformat(),
                'query': query[:100] + '...' if len(query) > 100 else query,
                'execution_time': f"{execution_time:.2f}s",
                'summary': {
                    'overview': f"执行了用户查询：{query[:50]}...",
                    'steps': f"总共执行了{total_steps}个步骤，包括：" + ", ".join([f"{k}({v}次)" for k, v in step_types.items()]),
                    'status': '执行完成',
                    'findings': '详细结果请查看执行过程',
                    'recommendations': '任务已完成'
                },
                'method': 'basic'
            }
            
        except Exception as e:
            logger.error(f"生成基础总结失败: {str(e)}")
            return {
                'type': 'error_summary',
                'timestamp': datetime.now().isoformat(),
                'query': query,
                'execution_time': f"{execution_time:.2f}s",
                'summary': {
                    'overview': '总结生成失败',
                    'status': '错误',
                    'error': str(e)
                },
                'method': 'error'
            }

# 创建全局实例
summary_service = SummaryService()